# Step1： Load documents
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders import Docx2txtLoader

base_dir = './python/langchain/OneFlower/docs'
documents = []

for file in os.listdir(base_dir):
    file_path = os.path.join(base_dir, file)
    if file.endswith('.pdf'):
        loader = PyPDFLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith('.docx') or file.endswith('.docx'):
       loader = Docx2txtLoader(file_path)
       documents.extend(loader.load())
    elif file.endswith('.txt'):
        loader = TextLoader(file_path)
        documents.extend(loader.load())

# Step 2: Splitting
from langchain.text_splitter import RecursiveCharacterTextSplitter

if not documents:
    print("错误: 没有从./docs目录加载到任何文档。请检查目录是否存在且包含支持的文件(PDF、DOCX、TXT)。")
    exit(1)

# 设置文档块大小为 200 字符
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=100)
chunked_documents = text_splitter.split_documents(documents)

# 检查文档分割是否成功
if not chunked_documents:
    print("错误: 文档分割后没有生成任何文档块。可能是文档内容太短或分割参数不合适。")
    exit(1)

print(f"成功加载并分割了 {len(chunked_documents)} 个文档块")

# Step3: Storage，word embedding to vectorstore
from langchain_community.vectorstores import Qdrant 
from langchain_ollama.embeddings import OllamaEmbeddings

try:
    vectorstore = Qdrant.from_documents(
      documents=chunked_documents,
      embedding = OllamaEmbeddings(
        model="nomic-embed-text",  # ollama embedding model
        base_url="http://localhost:11434",  # Ollama 的服务地址
      ),
      location=":memory:", # in-memory 存储， 使用内存存储模式，不需要连接外部Qdrant服务器
      collection_name="my_documents"
    )
    print("成功创建了向量存储")
except Exception as e:
    print(f"创建向量存储时出错: {str(e)}")
    exit(1)

# Step4: Retrieval VectorStore & LLMs
import logging
from langchain_ollama import ChatOllama
from langchain.retrievers.multi_query import MultiQueryRetriever # 多查询检索器
from langchain.chains import RetrievalQA # 检索问答链

logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

llm = ChatOllama(
    model="deepseek-r1:8b",  # 模型名称
    base_url="http://localhost:11434",  # Ollama 服务地址
    temperature=0.7,  # 创造性程度
    num_predict=512,  # 最大生成长度
)

# MultiQueryRetriever 工作原理
# 多查询生成: 使用LLM从原始问题生成多个相关问题
# 并行检索: 对所有生成的问题进行向量检索
# 结果去重: 合并并去重所有检索结果
# 返回最相关文档
retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(),
    llm=llm # 基于用户提问生成多个相关问题
)

# 检索问答链，检索增强生成（RAG）管道
qa_chain = RetrievalQA.from_chain_type(
    llm=llm, # 基于检索内容总结提炼
    retriever=retriever_from_llm, # 检索器，找到具体的‘嵌入片‘
    return_source_documents=True
)

# Step5：Output：生成回答并展示
from flask import Flask, request, render_template

app = Flask(__name__)

@app.route('/', methods=['GET', 'POST'])
def home():
    if request.method == 'POST':
      question = request.form.get('question')
      result = qa_chain({'query': question})
      return render_template('index.html', question=question, result=result)

    return render_template('index.html')

if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=True, port=5001)